import math
import tkinter
# import matplotlib
# matplotlib.use('TkAgg')
import os
import warnings

os.environ["OMP_NUM_THREADS"] = "8"
from inspect import currentframe, getframeinfo
import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn as nn
import scipy.io as scio
from CR_DSPPytorch import EDCLayer, PhaseRecLayer, FIRLayer, format_rt, cmul
import util

os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
BASE_DIR = os.path.dirname(__file__)


def TDC(lp_range=[1], tap_list=[2048], init_method='LS_CO', save_result_flag=False, clusters_flag=False,
        quantization_flag=False):
    frame = currentframe()
    experiment_name = getframeinfo(frame).function
    pick_syms_num = 131072
    modOrder = 4
    Qcache = np.zeros([len(lp_range), len(tap_list)])
    BERcache = np.zeros_like(Qcache)
    result_save_dir = os.path.join(BASE_DIR, f'result/TDE')
    if not os.path.exists(result_save_dir):
        os.makedirs(result_save_dir)
    chosen_device = 'cuda' if torch.cuda.is_available() else 'cpu'
    init_method = init_method
    save_result_flag = save_result_flag
    quantization_flag = quantization_flag

    for lpIndx, lp in enumerate(lp_range):
        for tapIndx, tap in enumerate(tap_list):
            # load data
            dataPath = os.path.join(BASE_DIR, 'data/simulation/16QAM32Gbaud3200kmHe/PMD_eliminated_by_CMA'
                                              f'/trSet_lp_{lp}.mat')
            data = scio.loadmat(dataPath)
            symbolRate = data['symbolRate'][0, 0]
            spanNum = data['spanNum'][0, 0]
            spanLen = data['spanLen'][0, 0]
            sig = data['sig']
            prbs = data['prbs']
            L = spanNum * spanLen
            D = 17e-6
            DL = D * L / 2

            sig = torch.from_numpy(sig[np.newaxis, ...])
            constellations = torch.from_numpy(util.CONST_16QAM)

            '''均衡色散所需总抽头书上界'''
            total_taps_limit = util.cal_tap_num_odd(DL=DL,
                                                    sample_freq=2 * symbolRate)
            print(f'taps upper bound: {total_taps_limit}')

            '''DSP procedure'''
            edc = EDCLayer(symbol_rate=symbolRate, DL=DL, case_num=2, tap=tap, sample_factor=2,
                           power_norm=True, init_method=init_method, time_domain=True)
            # scio.savemat(os.path.join(result_save_dir, f'filter_taps_{lp}_dBm_{tap}_taps.mat'),
            #              {'h': edc.h.data[0, ...].cpu().data.numpy()})
            # filter_tap_moduli(edc.h.data,
            #                   show=True,
            #                   saveFig=False,
            #                   savePath=os.path.join(result_save_dir, f'Moduli_of_filter_taps_init_by_{init_method}.png'),
            #                   splitRi=False)
            # heat_map_of_filter_taps(edc.h.data,
            #                         show=True,
            #                         saveFig=True,
            #                         savePath=os.path.join(result_save_dir, f'Heat_map_of_filter_taps_init_by_{init_method}.png'))
            # cluster filter coefficients
            if clusters_flag:
                edc.h.data = modified_filter_taps_by_cluster_complex(edc.h.data, clusters=16)
                # edc.h.data = modified_filter_taps_by_cluster_real(edc.h.data, clusters=8)
                # heat_map_of_filter_taps(edc.h.data,
                #                         show=True,
                #                         saveFig=True,
                #                         savePath=os.path.join(result_save_dir, f'Heat_map_of_filter_taps_init_by_{init_method}_clustered.png'))
                # filter_tap_moduli(edc.h.data,
                #                   show=True,
                #                   saveFig=True,
                #                   savePath=os.path.join(result_save_dir, f'Moduli_of_filter_taps_init_by_{init_method}_clustered.png'))
            if quantization_flag:
                edc.h.data = quantization_filter_taps(edc.h.data, bit_width=4)
            pr = PhaseRecLayer(1024)
            lms = FIRLayer(tap=32, case_num=2, power_norm=True, centor_one=True)

            sig = sig.to(chosen_device)
            lms = lms.to(chosen_device)
            edc = edc.to(chosen_device)

            sig = sig[..., 0:pick_syms_num * 2]
            prbs = prbs[..., 0:pick_syms_num * modOrder]

            sig = edc(torch.view_as_real(sig))
            sig = pr(sig[..., 1::2, :])
            lms.fit(sig, err_mode='DDM', constellations=constellations, iter_num=4, block_size=4028,
                    remain=2048, lr=5e-4)
            sig = lms(sig)

            sig = torch.view_as_complex(sig)
            sig = sig.cpu().data.numpy().squeeze()
            sig, ber, _ = util.pr_ber(sig, prbs, constellations.cpu().data.numpy())

            ber = np.mean(ber)
            BERcache[lpIndx, tapIndx] = ber
            Qcache[lpIndx, tapIndx] = util.ber2q(ber)
            print(f'lp: {lp}, tap:{tap}, BER: {ber}, Q: {util.ber2q(ber)}')

    if save_result_flag:
        with open(os.path.join(result_save_dir, f'TDE_vary_tap.csv'), 'w') as f:
            f.write('lp, tap, BER, Q\n')
            for lpIndx, lp in enumerate(lp_range):
                for tapIndx, tap in enumerate(tap_list):
                    f.write(f'{lp}, {tap}, {BERcache[lpIndx, tapIndx]}, {Qcache[lpIndx, tapIndx]}\n')
        print(f'result saved at \"{result_save_dir}\"')


def TDC_Cal_Complexity(lp=2, tap_list=[2048], cluster_list=[16], init_method='LS_CO', save_result_flag=False):
    frame = currentframe()
    experiment_name = getframeinfo(frame).function
    pick_syms_num = 131072
    modOrder = 4
    Qcache = np.zeros([len(tap_list), len(cluster_list)])
    BERcache = np.zeros_like(Qcache)
    Complexicache = np.zeros_like(Qcache)
    result_save_dir = os.path.join(BASE_DIR, f'result/TDE')
    if not os.path.exists(result_save_dir):
        os.makedirs(result_save_dir)
    chosen_device = 'cuda' if torch.cuda.is_available() else 'cpu'
    init_method = init_method
    save_result_flag = save_result_flag

    for tapIndx, tap in enumerate(tap_list):
        for clusterIndx, cluster in enumerate(cluster_list):
            # load data
            dataPath = os.path.join(BASE_DIR, 'data/simulation/16QAM32Gbaud3200kmHe/PMD_eliminated_by_CMA'
                                              f'/trSet_lp_{lp}.mat')
            data = scio.loadmat(dataPath)
            symbolRate = data['symbolRate'][0, 0]
            spanNum = data['spanNum'][0, 0]
            spanLen = data['spanLen'][0, 0]
            sig = data['sig']
            prbs = data['prbs']
            L = spanNum * spanLen
            D = 17e-6
            DL = D * L / 2

            sig = torch.from_numpy(sig[np.newaxis, ...])
            constellations = torch.from_numpy(util.CONST_16QAM)

            '''均衡色散所需总抽头书上界'''
            total_taps_limit = util.cal_tap_num_odd(DL=DL,
                                                    sample_freq=2 * symbolRate)
            print(f'taps upper bound: {total_taps_limit}')

            '''DSP procedure'''
            edc = EDCLayer(symbol_rate=symbolRate, DL=DL, case_num=2, tap=tap, sample_factor=2,
                           power_norm=True, init_method=init_method, time_domain=True)
            # filter_tap_moduli(edc.h.data,
            #                   show=True,
            #                   saveFig=False,
            #                   savePath=os.path.join(result_save_dir, f'Moduli_of_filter_taps_init_by_{init_method}.png'),
            #                   splitRi=False)
            # heat_map_of_filter_taps(edc.h.data,
            #                         show=True,
            #                         saveFig=True,
            #                         savePath=os.path.join(result_save_dir, f'Heat_map_of_filter_taps_init_by_{init_method}.png'))
            # cluster filter coefficients
            # edc.h.data = modified_filter_taps_by_cluster_complex(edc.h.data, clusters=16)
            edc.h.data = modified_filter_taps_by_cluster_complex(edc.h.data, clusters=cluster)
            # heat_map_of_filter_taps(edc.h.data,
            #                         show=True,
            #                         saveFig=True,
            #                         savePath=os.path.join(result_save_dir, f'Heat_map_of_filter_taps_init_by_{init_method}_clustered.png'))
            # filter_tap_moduli(edc.h.data,
            #                   show=True,
            #                   saveFig=True,
            #                   savePath=os.path.join(result_save_dir, f'Moduli_of_filter_taps_init_by_{init_method}_clustered.png'))
            pr = PhaseRecLayer(1024)
            lms = FIRLayer(tap=32, case_num=2, power_norm=True, centor_one=True)

            sig = sig.to(chosen_device)
            lms = lms.to(chosen_device)
            edc = edc.to(chosen_device)

            sig = sig[..., 0:pick_syms_num * 2]
            prbs = prbs[..., 0:pick_syms_num * modOrder]

            sig = edc(torch.view_as_real(sig))
            sig = pr(sig[..., 1::2, :])
            lms.fit(sig, err_mode='DDM', constellations=constellations, iter_num=4, block_size=4028,
                    remain=2048, lr=5e-4)
            sig = lms(sig)

            sig = torch.view_as_complex(sig)
            sig = sig.cpu().data.numpy().squeeze()
            sig, ber, _ = util.pr_ber(sig, prbs, constellations.cpu().data.numpy())
            Q = util.ber2q(np.mean(ber))
            BERcache[tapIndx, clusterIndx] = ber
            Qcache[tapIndx, clusterIndx] = Q
            Complexicache[tapIndx, clusterIndx] = 4 * cluster
            print(f'lp: {lp}, tap:{tap}, BER: {ber}, Q: {util.ber2q(ber)}')

    if save_result_flag:
        with open(os.path.join(result_save_dir, f'TDE_vary_tap_and_cluster.csv'), 'w') as f:
            f.write('tap, cluster, BER, Q, complexity\n')
            for tapIndx, tap in enumerate(tap_list):
                for clusterIndx, cluster in enumerate(cluster_list):
                    f.write(f' {tap}, {cluster}, {BERcache[tapIndx, clusterIndx]}, {Qcache[tapIndx, clusterIndx]}, '
                            f'{Complexicache[tapIndx, clusterIndx]}\n')
        print(f'result saved at \"{result_save_dir}\"')


def cluster_coeffs_of_TDC(lp_range=[1], tap=2048, cluster_list=[4], init_method='LS_CO', save_result_flag=False,
                          clusters_flag=False,
                          quantization_flag=False):
    frame = currentframe()
    experiment_name = getframeinfo(frame).function
    pick_syms_num = 131072
    modOrder = 4
    Qcache = np.zeros([len(lp_range), len(cluster_list)])
    BERcache = np.zeros_like(Qcache)
    result_save_dir = os.path.join(BASE_DIR, f'result/{experiment_name}')
    if not os.path.exists(result_save_dir):
        os.makedirs(result_save_dir)
    chosen_device = 'cuda' if torch.cuda.is_available() else 'cpu'
    init_method = init_method
    save_result_flag = save_result_flag
    quantization_flag = quantization_flag

    for lpIndx, lp in enumerate(lp_range):
        for clusterIndx, cluster in enumerate(cluster_list):
            print(f'lp={lp}, tap={tap} cluster={cluster}')
            # load data
            dataPath = os.path.join(BASE_DIR, 'data/simulation/16QAM32Gbaud3200kmHe/PMD_eliminated_by_CMA'
                                              f'/trSet_lp_{lp}.mat')
            data = scio.loadmat(dataPath)
            symbolRate = data['symbolRate'][0, 0]
            spanNum = data['spanNum'][0, 0]
            spanLen = data['spanLen'][0, 0]
            sig = data['sig']
            prbs = data['prbs']
            L = spanNum * spanLen
            D = 17e-6
            DL = D * L / 2

            sig = torch.from_numpy(sig[np.newaxis, ...])
            constellations = torch.from_numpy(util.CONST_16QAM)

            '''均衡色散所需总抽头书上界'''
            total_taps_limit = util.cal_tap_num_odd(DL=DL,
                                                    sample_freq=2 * symbolRate)
            print(f'taps upper bound: {total_taps_limit}')

            '''DSP procedure'''
            edc = EDCLayer(symbol_rate=symbolRate, DL=DL, case_num=2, tap=tap, sample_factor=2,
                           power_norm=True, init_method=init_method, time_domain=True)
            # filter_tap_moduli(edc.h.data,
            #                   show=True,
            #                   saveFig=True,
            #                   savePath=os.path.join(result_save_dir, f'Moduli_of_filter_taps_init_by_{init_method}.png'))
            # heat_map_of_filter_taps(edc.h.data,
            #                         show=True,
            #                         saveFig=True,
            #                         savePath=os.path.join(result_save_dir, f'Heat_map_of_filter_taps_init_by_{init_method}.png'))
            # cluster filter coefficients
            if clusters_flag:
                # edc.h.data = modified_filter_taps_by_cluster_real(edc.h.data, clusters=cluster)
                edc.h.data = modified_filter_taps_by_cluster_complex(edc.h.data, clusters=cluster)
                # heat_map_of_filter_taps(edc.h.data,
                #                         show=True)
                # filter_tap_moduli(edc.h.data,
                #                   show=True)
            if quantization_flag:
                edc.h.data = quantization_filter_taps(edc.h.data, bit_width=4)
            pr = PhaseRecLayer(1024)
            lms = FIRLayer(tap=32, case_num=2, power_norm=True, centor_one=True)

            sig = sig.to(chosen_device)
            lms = lms.to(chosen_device)
            edc = edc.to(chosen_device)

            sig = sig[..., 0:pick_syms_num * 2]
            prbs = prbs[..., 0:pick_syms_num * modOrder]

            sig = edc(torch.view_as_real(sig))
            sig = pr(sig[..., 1::2, :])
            lms.fit(sig, err_mode='DDM', constellations=constellations, iter_num=4, block_size=4028,
                    remain=2048, lr=5e-4)
            sig = lms(sig)

            sig = torch.view_as_complex(sig)
            sig = sig.cpu().data.numpy().squeeze()
            sig, ber, _ = util.pr_ber(sig, prbs, constellations.cpu().data.numpy())

            ber = np.mean(ber)
            BERcache[lpIndx, clusterIndx] = ber
            Qcache[lpIndx, clusterIndx] = util.ber2q(ber)
            print(f'lp: {lp}, tap:{tap}, cluster: {cluster}, BER: {ber}, Q: {util.ber2q(ber)}')

    if save_result_flag:
        with open(os.path.join(result_save_dir, f'Complex_TDC_init_{init_method}_vary_clusters_12_15.csv'), 'w') as f:
            f.write('lp, cluster, tap, BER, Q\n')
            for lpIndx, lp in enumerate(lp_range):
                for clusterIndx, cluster in enumerate(cluster_list):
                    f.write(f'{lp}, {cluster}, {tap}, {BERcache[lpIndx, clusterIndx]}, {Qcache[lpIndx, clusterIndx]}\n')
        print(f'result saved at \"{result_save_dir}\"')


def cluster_coeffs_of_clusterd_TDC(lp_range=[1], tap=2048, bit_width_list=[4], init_method='LS_CO',
                                   save_result_flag=False,
                                   clusters_flag=False,
                                   quantization_flag=False):
    frame = currentframe()
    experiment_name = getframeinfo(frame).function
    pick_syms_num = 131072
    modOrder = 4
    Qcache = np.zeros([len(lp_range), len(bit_width_list)])
    BERcache = np.zeros_like(Qcache)
    result_save_dir = os.path.join(BASE_DIR, f'result/{experiment_name}')
    if not os.path.exists(result_save_dir):
        os.makedirs(result_save_dir)
    chosen_device = 'cuda' if torch.cuda.is_available() else 'cpu'
    init_method = init_method
    cluster = 14
    save_result_flag = save_result_flag
    quantization_flag = quantization_flag

    for lpIndx, lp in enumerate(lp_range):
        for bitIndx, bit_width in enumerate(bit_width_list):
            print(f'lp={lp}, tap={tap} bit_width={bit_width}')
            # load data
            dataPath = os.path.join(BASE_DIR, 'data/simulation/16QAM32Gbaud3200kmHe/PMD_eliminated_by_CMA'
                                              f'/trSet_lp_{lp}.mat')
            data = scio.loadmat(dataPath)
            symbolRate = data['symbolRate'][0, 0]
            spanNum = data['spanNum'][0, 0]
            spanLen = data['spanLen'][0, 0]
            sig = data['sig']
            prbs = data['prbs']
            L = spanNum * spanLen
            D = 17e-6
            DL = D * L / 2

            sig = torch.from_numpy(sig[np.newaxis, ...])
            constellations = torch.from_numpy(util.CONST_16QAM)

            '''均衡色散所需总抽头书上界'''
            total_taps_limit = util.cal_tap_num_odd(DL=DL,
                                                    sample_freq=2 * symbolRate)
            print(f'taps upper bound: {total_taps_limit}')

            '''DSP procedure'''
            edc = EDCLayer(symbol_rate=symbolRate, DL=DL, case_num=2, tap=tap, sample_factor=2,
                           power_norm=True, init_method=init_method, time_domain=True)
            # filter_tap_moduli(edc.h.data,
            #                   show=True,
            #                   saveFig=True,
            #                   savePath=os.path.join(result_save_dir, f'Moduli_of_filter_taps_init_by_{init_method}.png'))
            # heat_map_of_filter_taps(edc.h.data,
            #                         show=True,
            #                         saveFig=True,
            #                         savePath=os.path.join(result_save_dir, f'Heat_map_of_filter_taps_init_by_{init_method}.png'))
            # cluster filter coefficients
            if clusters_flag:
                edc.h.data = modified_filter_taps_by_cluster_real(edc.h.data, clusters=cluster)
                # edc.h.data = modified_filter_taps_by_cluster_complex(edc.h.data, clusters=14)
                # heat_map_of_filter_taps(edc.h.data,
                #                         show=True)
                # filter_tap_moduli(edc.h.data,
                #                   show=True)
            if quantization_flag:
                edc.h.data = quantization_filter_taps(edc.h.data, bit_width=bit_width)
            pr = PhaseRecLayer(1024)
            lms = FIRLayer(tap=32, case_num=2, power_norm=True, centor_one=True)

            sig = sig.to(chosen_device)
            lms = lms.to(chosen_device)
            edc = edc.to(chosen_device)

            sig = sig[..., 0:pick_syms_num * 2]
            prbs = prbs[..., 0:pick_syms_num * modOrder]

            sig = edc(torch.view_as_real(sig))
            sig = pr(sig[..., 1::2, :])
            lms.fit(sig, err_mode='DDM', constellations=constellations, iter_num=4, block_size=4028,
                    remain=2048, lr=5e-4)
            sig = lms(sig)

            sig = torch.view_as_complex(sig)
            sig = sig.cpu().data.numpy().squeeze()
            sig, ber, _ = util.pr_ber(sig, prbs, constellations.cpu().data.numpy())

            ber = np.mean(ber)
            BERcache[lpIndx, bitIndx] = ber
            Qcache[lpIndx, bitIndx] = util.ber2q(ber)
            print(f'lp: {lp},bit_width: {bit_width}, tap:{tap}, cluster: {cluster}, BER: {ber}, Q: {util.ber2q(ber)}')

    if save_result_flag:
        with open(os.path.join(result_save_dir, f'TDC_init_{init_method}_vary_bits_and_lp.csv'), 'w') as f:
            f.write('lp, bit_width, cluster, tap, BER, Q\n')
            for lpIndx, lp in enumerate(lp_range):
                for bitIndx, bit_width in enumerate(bit_width_list):
                    f.write(
                        f'{lp}, {bit_width}, {cluster}, {tap}, {BERcache[lpIndx, bitIndx]}, {Qcache[lpIndx, bitIndx]}\n')
        print(f'result saved at \"{result_save_dir}\"')


def quantize_coeffs_of_TDC(lp=1, tap=2048, init_method='LS_CO', save_result_flag=False,
                           clusters_flag=False, quantization_flag=False):
    frame = currentframe()
    experiment_name = getframeinfo(frame).function
    pick_syms_num = 131072
    modOrder = 4
    result_save_dir = os.path.join(BASE_DIR, f'result/{experiment_name}')
    if not os.path.exists(result_save_dir):
        os.makedirs(result_save_dir)
    chosen_device = 'cuda' if torch.cuda.is_available() else 'cpu'
    init_method = init_method
    save_result_flag = save_result_flag
    quantization_flag = quantization_flag
    bit_width_list = [2, 3, 4, 5, 6, 7, 8]
    BERcache = np.zeros([len(bit_width_list)])
    Qcache = np.zeros_like(BERcache)
    for bitIndx, bit_width in enumerate(bit_width_list):
        print(f'bit width: {bit_width}')
        # load data
        dataPath = os.path.join(BASE_DIR, 'data/simulation/16QAM32Gbaud3200kmHe/PMD_eliminated_by_CMA'
                                          f'/trSet_lp_{lp}.mat')
        data = scio.loadmat(dataPath)
        symbolRate = data['symbolRate'][0, 0]
        spanNum = data['spanNum'][0, 0]
        spanLen = data['spanLen'][0, 0]
        sig = data['sig']
        prbs = data['prbs']
        L = spanNum * spanLen
        D = 17e-6
        DL = D * L / 2

        sig = torch.from_numpy(sig[np.newaxis, ...])
        constellations = torch.from_numpy(util.CONST_16QAM)

        '''均衡色散所需总抽头书上界'''
        total_taps_limit = util.cal_tap_num_odd(DL=DL,
                                                sample_freq=2 * symbolRate)
        print(f'taps upper bound: {total_taps_limit}')

        '''DSP procedure'''
        edc = EDCLayer(symbol_rate=symbolRate, DL=DL, case_num=2, tap=tap, sample_factor=2,
                       power_norm=True, init_method=init_method, time_domain=True)
        # filter_tap_moduli(edc.h.data,
        #                   show=True,
        #                   saveFig=True,
        #                   savePath=os.path.join(result_save_dir, f'Moduli_of_filter_taps_init_by_{init_method}.png'))
        # heat_map_of_filter_taps(edc.h.data,
        #                         show=True,
        #                         saveFig=True,
        #                         savePath=os.path.join(result_save_dir, f'Heat_map_of_filter_taps_init_by_{init_method}.png'))
        # cluster filter coefficients
        if clusters_flag:
            # edc.h.data = modified_filter_taps_by_cluster_in_complex_plane(edc.h.data, clusters=14)
            edc.h.data = modified_filter_taps_by_cluster_real(edc.h.data, clusters=14)
            # heat_map_of_filter_taps(edc.h.data,
            #                         show=True)
            # filter_tap_moduli(edc.h.data,
            #                   show=True)
        if quantization_flag:
            edc.h.data = quantization_filter_taps(edc.h.data, bit_width=bit_width)
        pr = PhaseRecLayer(1024)
        lms = FIRLayer(tap=32, case_num=2, power_norm=True, centor_one=True)

        sig = sig.to(chosen_device)
        lms = lms.to(chosen_device)
        edc = edc.to(chosen_device)

        sig = sig[..., 0:pick_syms_num * 2]
        prbs = prbs[..., 0:pick_syms_num * modOrder]

        sig = edc(torch.view_as_real(sig))
        sig = pr(sig[..., 1::2, :])
        lms.fit(sig, err_mode='DDM', constellations=constellations, iter_num=4, block_size=4028,
                remain=2048, lr=5e-4)
        sig = lms(sig)

        sig = torch.view_as_complex(sig)
        sig = sig.cpu().data.numpy().squeeze()
        sig, ber, _ = util.pr_ber(sig, prbs, constellations.cpu().data.numpy())

        ber = np.mean(ber)
        BERcache[bitIndx] = ber
        Qcache[bitIndx] = util.ber2q(ber)
        print(f'lp: {lp}, tap:{tap}, BER: {ber}, bit_width: {bit_width} Q: {util.ber2q(ber)}')

    if save_result_flag:
        with open(os.path.join(result_save_dir, f'Quantize_coeffs_of_TDC_init_{init_method}.csv'), 'w') as f:
            f.write('bit_width, lp, tap, BER, Q\n')
            for bitIndx, bit_width in enumerate(bit_width_list):
                f.write(f'{bit_width}, {lp}, {tap}, {BERcache[bitIndx]}, {Qcache[bitIndx]}\n')
        print(f'result saved at \"{result_save_dir}\"')


def filter_tap_moduli(h, show=False, saveFig=False, savePath=None, splitRi=False):
    '''
    h: filter taps: [2, tap_num, re/im], torch.Size
    plot the moduli of taps
    '''
    if saveFig == True:
        if savePath is None:
            raise Exception('savePath is supposed to be assigned when saveFig is True.')
    # 计算滤波器抽头的模值
    h = h.cpu().data.numpy().astype(np.float32)
    if not splitRi:
        h_moduli = np.sqrt(h[..., 0] ** 2 + h[..., 1] ** 2)
        h_moduli = h_moduli[0, :]
        plt.plot(h_moduli, 'o-', markersize=4.5)
        plt.xlabel('Position')
        plt.ylabel('Moduli')
        plt.grid(True)
    else:
        h_moduli_r = h[0, ..., 0]
        h_moduli_i = h[0, ..., 1]
        fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(8, 6))
        ax1.plot(h_moduli_r, 'o-', markersize=4.5, label='Real Part')
        ax2.plot(h_moduli_i, 'o-', markersize=4.5, label='Imaginary Part')
        plt.tight_layout()

    if saveFig:
        plt.savefig(savePath, dpi=300)
    if show:
        plt.show()
    plt.close()


def heat_map_of_filter_taps(h, show=False, saveFig=False, savePath=None):
    '''
    h: filter taps: [2, tap_num, re/im], torch.Size
    plot the heat map of taps
    '''
    if saveFig == True:
        if savePath is None:
            raise Exception('savePath is supposed to be assigned when saveFig is True.')
    h = h.cpu().data.numpy().astype(np.float32)
    h = h[0, ...]
    h_real = h[..., 0]
    h_imag = h[..., 1]
    h_magn = np.sqrt(h_real ** 2 + h_imag ** 2)
    h_phase = np.arctan2(h_imag, h_real)
    # 根据相位绘制热力图
    plt.scatter(h_real.flatten(), h_imag.flatten(), c=h_phase.flatten(), cmap='viridis', s=100)
    # plt.hist2d(h_real.flatten(), h_imag.flatten(), bins=50, cmap='viridis')
    plt.title('Heatmap of Filter Taps in the Complex Plane (phase)')
    plt.xlabel('Filter Coefficients (Real)')
    plt.ylabel('Filter Coefficients (Imaginary)')
    plt.grid(False)
    if saveFig:
        plt.savefig(savePath, dpi=300)
    if show:
        plt.show()
    plt.close()


def modified_filter_taps_by_cluster_complex(h_ori, clusters=8):
    '''
    h: filter taps: [2, tap_num, re/im], torch.Size
    return: modified filter taps
    Implement clustering in the complex plane
    '''
    from sklearn.cluster import KMeans
    h_ori = h_ori.cpu().data.numpy().astype(np.float32)
    # 转化为2D数组
    h = h_ori.reshape(-1, 2)
    kmeans = KMeans(n_clusters=clusters, random_state=42, n_init=10)
    kmeans.fit(h)
    # 获取聚类中心和每个系数的聚类标签
    centroids = kmeans.cluster_centers_
    labels = kmeans.labels_
    new_h = centroids[labels]
    new_h = new_h.reshape(h_ori.shape)
    return torch.tensor(new_h, dtype=torch.float32)


def modified_filter_taps_by_cluster_real(h_ori, clusters=8):
    '''
    h: filter taps: [2, tap_num, re/im], torch.Size
    return: modified filter taps
    '''
    from sklearn.cluster import KMeans
    h_ori = h_ori.cpu().data.numpy().astype(np.float64)
    # 转化为2D数组
    h = h_ori.reshape(-1, 2)

    real_part = h[..., 0][..., np.newaxis]
    kmeas_real = KMeans(n_clusters=clusters, random_state=42, n_init=10)
    kmeas_real.fit(real_part)
    centroids_real = kmeas_real.cluster_centers_
    labels_real = kmeas_real.labels_
    new_real_part = centroids_real[labels_real]

    imag_part = h[..., 1][..., np.newaxis]
    kmeas_imag = KMeans(n_clusters=clusters, random_state=42, n_init=10)
    kmeas_imag.fit(imag_part)
    centroids_imag = kmeas_imag.cluster_centers_
    labels_imag = kmeas_imag.labels_
    new_imag_part = centroids_imag[labels_imag]

    new_h = np.hstack([new_real_part, new_imag_part])
    new_h = new_h.reshape(h_ori.shape)
    return torch.tensor(new_h, dtype=torch.float32)


def quantization_filter_taps(h, bit_width=8):
    '''
    h: filter taps: [2, tap_num, re/im], torch.Size
    bit_width: bit width of quantization
    return: quantized filter taps
    '''
    from CR_DSPPytorch import cal_uniform_quantization_parameters, uniform_quant
    sx_r, ax_r = cal_uniform_quantization_parameters(
        h[..., 0], bit_width=bit_width, axis=-1, approach='maxabs')
    sx_i, ax_i = cal_uniform_quantization_parameters(
        h[..., 1], bit_width=bit_width, axis=-1, approach='maxabs')
    h_r = uniform_quant((h[..., 0], sx_r, ax_r), bit_width=bit_width, axis=-2, approach='maxabs')
    h_i = uniform_quant((h[..., 1], sx_i, ax_i), bit_width=bit_width, axis=-2, approach='maxabs')
    new_h = torch.stack([h_r, h_i], dim=-1)
    assert new_h.shape == h.shape
    return new_h


if __name__ == '__main__':
    warnings.filterwarnings("ignore", category=UserWarning)
    lp_range = list(np.arange(-2, 4, 1))
    # tap_list = [2048, 1024, 512]
    init_method = 'TDS'
    tap_list = [427, 447, 507, 557, 607, 657, 707]
    TDC(lp_range=[2], tap_list=[597], init_method=init_method,
        save_result_flag=False, clusters_flag=True, quantization_flag=False)
    # cluster_coeffs_of_TDC(lp_range=[1], tap=897, cluster_list=[8, 16, 32, 48, 64, 80], init_method=init_method,
    #                       save_result_flag=True, clusters_flag=True, quantization_flag=False)
    # cluster_coeffs_of_clusterd_TDC(lp_range=lp_range, tap=597, bit_width_list=[2, 3, 4],
    #                                init_method=init_method,
    #                                save_result_flag=True, clusters_flag=True, quantization_flag=True)
    # quantize_coeffs_of_TDC(lp=0, tap=597, init_method=init_method,
    #                        save_result_flag=True, clusters_flag=False, quantization_flag=True)
